
import os.path
import struct
import mmap

from archive_base import Entry

class NSAEntry(Entry):
	def __init__(self, name, offset, size):
		self.name = name
		self.offset = offset
		self.size = size
	def decode(self):
		with open(self.path, 'rb') as stream:
			stream.seek(self.offset)
			return stream.read(self.size)

def read_cstring(stream):
	bytes = []
	while True:
		byte = stream.read(1)
		if byte == '\x00':
			break
		elif byte == '':
			raise EOFError()
		else:
			bytes.append(byte)
	return ''.join(bytes)

def m_read_cstring(stream):
	z = stream.find('\x00')
	text = stream[stream.tell():z]
	stream.seek(z+1)
	return text

def read_indexes_internal(stream, read_cstring):
	n = struct.unpack('>H', stream.read(2))[0]
	global_offset = struct.unpack('>I', stream.read(4))[0]

	indexes = []
	for i in range(n):
		name = unicode(read_cstring(stream)) # XXX: don't support non-ascii
		offset, size, expandsize = struct.unpack('>xIII', stream.read(1+4+4+4))
		assert size == expandsize
		indexes.append(NSAEntry(name, global_offset + offset, size))
	return indexes

def read_indexes(path):
	with open(path, 'rb') as x:
		try:
			m = mmap.mmap(x.fileno(), 0, access=mmap.ACCESS_READ)
		except Exception, e:
			print 'ignoring error', e
			return read_indexes_internal(x, read_cstring)
		else:
			indexes = read_indexes_internal(m, m_read_cstring)
			m.close()
			return indexes

def extract_nsa_to(nsa, root):
	indexes = read_indexes(nsa)
	if not os.path.exists(root):
		os.makedirs(root)
	with open(nsa, 'rb') as r:
		for index in indexes:
			dir = os.path.join(root, os.path.dirname(index.name))
			if not os.path.exists(dir):
				os.makedirs(dir)
			path = os.path.join(root, index.name)
			with open(path, 'wb') as x:
				r.seek(index.offset)
				x.write(r.read(index.size))


if __name__ == '__main__':
	from time import time
	start = time()
	indexes = read_indexes('arc3.nsa')
	print time() - start
	#extract_nsa_to('arc1.nsa', 'x')



